package com.woniu.utils;

import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @Author: 朱杰
 * @Description: 用户上下文管理工具类
 * @DateTime: 2025/6/11 19:48
 **/
@Component
public class ConversationContextManager {

    private final Map<String, List<String>> conversationContexts = new HashMap<>();

    private static final String USER_PREFIX = "USER: ";
    private static final String AI_PREFIX = "AI: ";

    /**
     * 添加用户的提问到上下文记录中
     * @param sessionId 用户会话ID
     * @param userInput 用户的提问
     */
    public void addUserInput(String sessionId, String userInput) {
        conversationContexts.computeIfAbsent(sessionId, k -> new ArrayList<>()).add(USER_PREFIX + userInput);
    }

    /**
     * 添加AI的回复到上下文记录中
     * @param sessionId 用户会话ID
     * @param aiResponse AI的回复
     */
    public void addAIResponse(String sessionId, String aiResponse) {
        // 合并流式响应中的换行
        String mergedResponse = aiResponse.replaceAll("\n", " ");
        conversationContexts.computeIfAbsent(sessionId, k -> new ArrayList<>()).add(AI_PREFIX + mergedResponse);
    }

    /**
     * 获取用户的上下文记录
     * @param sessionId 用户会话ID
     * @return 上下文记录
     */
    public List<String> getContext(String sessionId) {
        return conversationContexts.getOrDefault(sessionId, new ArrayList<>());
    }

    /**
     * 清除用户的上下文记录
     * @param sessionId 用户会话ID
     */
    public void clearContext(String sessionId) {
        conversationContexts.remove(sessionId);
    }

}
